"""
人数统计系统主程序
整合YOLOv11检测、ByteTrack跟踪、区域计数等功能
"""

import cv2
import numpy as np
import yaml
import logging
import os
import time
import argparse
from typing import Optional, Dict, Any

# 导入自定义模块
from detector import PersonDetector, DetectionVisualizer
from tracker import PersonTracker, TrackingVisualizer
from region_manager import RegionManager, RegionConfig
from counter import PersonCounter, CounterVisualizer
from video_processor import RTSPProcessor, VideoSaver, FPSCounter

class PeopleCountingSystem:
    """人数统计系统主类"""
    
    def __init__(self, config_path: str = "config.yaml"):
        """
        初始化系统
        
        Args:
            config_path: 配置文件路径
        """
        self.config_path = config_path
        self.config = self._load_config()
        
        # 初始化日志
        self._setup_logging()
        
        # 初始化组件
        self.detector = None
        self.tracker = None
        self.region_manager = None
        self.counter = None
        self.video_processor = None
        self.video_saver = None
        
        # 可视化器
        self.detection_visualizer = DetectionVisualizer()
        self.tracking_visualizer = TrackingVisualizer()
        self.counter_visualizer = CounterVisualizer()
        
        # 性能监控
        self.fps_counter = FPSCounter()
        
        # 运行状态
        self.is_running = False
        self.setup_complete = False
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("人数统计系统初始化完成")
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            return config
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict[str, Any]:
        """获取默认配置"""
        return {
            'video': {
                'rtsp_url': 'rtsp://admin:Lorytech2015@192.168.3.79:554/cam/realmonitor?channel=1&subtype=0',
                'fps_limit': 15,
                'resize_width': 1280,
                'resize_height': 720
            },
            'model': {
                'model_path': 'yolo11n.pt',
                'confidence_threshold': 0.5,
                'iou_threshold': 0.45,
                'classes': [0]
            },
            'tracking': {
                'track_thresh': 0.25,
                'track_buffer': 30,
                'match_thresh': 0.8
            },
            'region': {
                'margin': {'top': 50, 'bottom': 50, 'left': 50, 'right': 50},
                'points': [],
                'color': [0, 255, 0],
                'thickness': 2
            },
            'display': {
                'show_video': True,
                'save_video': False,
                'output_path': 'output/result.mp4'
            }
        }
    
    def _setup_logging(self):
        """设置日志"""
        log_config = self.config.get('logging', {})
        level = getattr(logging, log_config.get('level', 'INFO'))
        
        # 创建日志目录
        log_file = log_config.get('file_path', 'logs/system.log')
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        # 配置日志
        logging.basicConfig(
            level=level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler() if log_config.get('console', True) else logging.NullHandler()
            ]
        )
    
    def setup_components(self):
        """设置系统组件"""
        try:
            # 初始化检测器
            model_config = self.config['model']
            self.detector = PersonDetector(
                model_path=model_config['model_path'],
                confidence_threshold=model_config['confidence_threshold'],
                iou_threshold=model_config['iou_threshold'],
                classes=model_config['classes']
            )
            
            # 初始化跟踪器
            tracking_config = self.config['tracking']
            self.tracker = PersonTracker(
                track_thresh=tracking_config['track_thresh'],
                track_buffer=tracking_config['track_buffer'],
                match_thresh=tracking_config['match_thresh'],
                frame_rate=self.config['video']['fps_limit']
            )
            
            # 初始化视频处理器
            video_config = self.config['video']
            resize_shape = None
            if video_config.get('resize_width') and video_config.get('resize_height'):
                resize_shape = (video_config['resize_width'], video_config['resize_height'])
            
            self.video_processor = RTSPProcessor(
                rtsp_url=video_config['rtsp_url'],
                fps_limit=video_config['fps_limit'],
                resize_shape=resize_shape
            )
            
            # 初始化视频保存器（如果需要）
            if self.config['display']['save_video']:
                output_path = self.config['display']['output_path']
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                
                frame_size = resize_shape or (1280, 720)
                self.video_saver = VideoSaver(
                    output_path=output_path,
                    fps=video_config['fps_limit'],
                    frame_size=frame_size
                )
            
            self.setup_complete = True
            self.logger.info("系统组件设置完成")
            return True
            
        except Exception as e:
            self.logger.error(f"设置系统组件失败: {e}")
            return False
    
    def setup_region(self) -> bool:
        """设置检测区域"""
        region_config = self.config['region']
        
        # 创建区域配置
        config = RegionConfig(
            points=region_config.get('points', []),
            margin=region_config.get('margin', {'top': 50, 'bottom': 50, 'left': 50, 'right': 50}),
            color=tuple(region_config.get('color', [0, 255, 0])),
            thickness=region_config.get('thickness', 2)
        )
        
        self.region_manager = RegionManager(config)
        
        # 如果没有预设区域点，需要交互式设置
        if not region_config.get('points'):
            self.logger.info("未找到预设区域，启动交互式区域设置")
            
            # 获取参考帧
            if not self.video_processor.start_capture():
                self.logger.error("无法启动视频捕获用于区域设置")
                return False
            
            # 等待获取帧
            reference_frame = None
            for _ in range(10):  # 尝试10次
                frame = self.video_processor.get_frame(timeout=2.0)
                if frame is not None:
                    reference_frame = frame
                    break
                time.sleep(0.1)
            
            if reference_frame is None:
                self.logger.error("无法获取参考帧")
                self.video_processor.stop_capture()
                return False
            
            # 交互式设置区域
            success = self.region_manager.interactive_region_setup(reference_frame)
            self.video_processor.stop_capture()
            
            if not success:
                self.logger.error("区域设置失败")
                return False
            
            # 保存区域配置
            self.region_manager.save_region_config("region_config.json")
            self.logger.info("区域配置已保存")
        
        # 初始化计数器
        self.counter = PersonCounter(self.region_manager)
        
        return True
    
    def run(self):
        """运行主程序"""
        if not self.setup_complete:
            self.logger.error("系统未正确设置")
            return
        
        if not self.setup_region():
            self.logger.error("区域设置失败")
            return
        
        # 启动视频处理
        if not self.video_processor.start_capture():
            self.logger.error("无法启动视频捕获")
            return
        
        # 启动视频保存（如果需要）
        if self.video_saver:
            self.video_saver.start_recording()
        
        self.is_running = True
        self.logger.info("系统开始运行")
        
        try:
            self._main_loop()
        except KeyboardInterrupt:
            self.logger.info("用户中断程序")
        except Exception as e:
            self.logger.error(f"运行时错误: {e}")
        finally:
            self._cleanup()
    
    def _main_loop(self):
        """主处理循环"""
        print("人数统计系统运行中...")
        print("按 'q' 键退出")
        print("按 'r' 键重置计数器")
        print("按 's' 键显示统计信息")
        
        while self.is_running:
            # 获取视频帧
            frame = self.video_processor.get_frame(timeout=1.0)
            if frame is None:
                continue
            
            # 处理帧
            processed_frame = self._process_frame(frame)
            
            # 显示结果
            if self.config['display']['show_video']:
                cv2.imshow("人数统计系统", processed_frame)
                
                # 处理按键
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    break
                elif key == ord('r'):
                    self.counter.reset_counter()
                    self.logger.info("计数器已重置")
                elif key == ord('s'):
                    self._print_statistics()
            
            # 保存视频帧
            if self.video_saver:
                self.video_saver.write_frame(processed_frame)
    
    def _process_frame(self, frame: np.ndarray) -> np.ndarray:
        """处理单帧"""
        # 检测人员
        detections = self.detector.detect(frame)
        
        # 跟踪人员
        tracked_detections = self.tracker.update(detections)
        track_infos = self.tracker.get_track_info(tracked_detections)
        
        # 更新计数器
        count_status = self.counter.update(track_infos)
        
        # 可视化
        annotated_frame = frame.copy()
        
        # 绘制区域
        annotated_frame = self.region_manager.draw_region(annotated_frame)
        
        # 绘制跟踪结果
        annotated_frame = self.tracking_visualizer.annotate_frame(
            annotated_frame, tracked_detections, show_traces=True
        )
        
        # 绘制计数信息
        annotated_frame = self.counter_visualizer.draw_counter_info(
            annotated_frame, count_status
        )
        
        # 绘制人员状态
        annotated_frame = self.counter_visualizer.draw_person_states(
            annotated_frame, track_infos, self.counter.person_states
        )
        
        # 绘制系统信息
        fps = self.fps_counter.update()
        video_stats = self.video_processor.get_stats()
        
        # 系统信息
        info_y = annotated_frame.shape[0] - 100
        cv2.putText(annotated_frame, f"FPS: {fps:.1f}", (10, info_y), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        cv2.putText(annotated_frame, f"Queue: {video_stats['queue_size']}", (10, info_y + 25), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        cv2.putText(annotated_frame, f"Connected: {video_stats['connected']}", (10, info_y + 50), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        return annotated_frame
    
    def _print_statistics(self):
        """打印统计信息"""
        count_status = self.counter.get_count_status()
        statistics = self.counter.get_statistics()
        video_stats = self.video_processor.get_stats()
        
        print("\n" + "="*50)
        print("系统统计信息")
        print("="*50)
        print(f"进入人数: {count_status['enter_count']}")
        print(f"离开人数: {count_status['exit_count']}")
        print(f"当前人数: {count_status['current_count']}")
        print(f"活跃跟踪: {count_status['active_tracks']}")
        print(f"实际FPS: {video_stats['fps_actual']:.1f}")
        print(f"总记录数: {statistics['total_records']}")
        print(f"运行时间: {statistics.get('time_span', 0):.1f}秒")
        print("="*50)
    
    def _cleanup(self):
        """清理资源"""
        self.is_running = False
        
        if self.video_processor:
            self.video_processor.stop_capture()
        
        if self.video_saver:
            self.video_saver.stop_recording()
        
        cv2.destroyAllWindows()
        
        self.logger.info("系统已停止，资源已清理")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="人数统计系统")
    parser.add_argument("--config", default="config.yaml", help="配置文件路径")
    parser.add_argument("--setup-region", action="store_true", help="只设置区域")
    args = parser.parse_args()
    
    # 创建系统实例
    system = PeopleCountingSystem(args.config)
    
    # 设置组件
    if not system.setup_components():
        print("系统设置失败")
        return
    
    # 如果只是设置区域
    if args.setup_region:
        if system.setup_region():
            print("区域设置完成")
        else:
            print("区域设置失败")
        return
    
    # 运行系统
    system.run()


if __name__ == "__main__":
    main() 